import base64

from Crypto import Random
from Crypto.Hash import SHA, SHA1
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5
from Crypto.Signature import PKCS1_v1_5 as Signature_pkcs1_v1_5
from Crypto.PublicKey import RSA

# 伪随机数生成器
from itsdangerous import signer

random_generator = Random.new().read
# rsa算法生成实例
rsa = RSA.generate(1024, random_generator)


def generate_server_key():
    # Server的秘钥对的生成
    private_pem = rsa.exportKey()
    with open("server-private.pem", "wb") as f:
        f.write(private_pem)

    public_pem = rsa.publickey().exportKey()
    with open("server-public.pem", "wb") as f:
        f.write(public_pem)


def generate_client_key():
    # Client的秘钥对的生成
    private_pem = rsa.exportKey()
    with open("client-private.pem", "wb") as f:
        f.write(private_pem)

    public_pem = rsa.publickey().exportKey()
    with open("client-public.pem", "wb") as f:
        f.write(public_pem)


def get_client_public_key():
    with open("client-public.pem") as f:
        return f.read()


def get_client_private_key():
    with open("client-private.pem") as f:
        return f.read()


def get_server_public_key():
    with open("server-public.pem") as f:
        return f.read()


def get_server_private_key():
    with open("server-private.pem") as f:
        return f.read()


def client_encrypt(message):
    # 用服务器的公匙加密
    with open("server-public.pem") as f:
        key = f.read()
        rsakey = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsakey)
        cipher_text = base64.b64encode(cipher.encrypt(message.encode('utf-8')))
        return cipher_text.decode('utf-8')


def client_decrypt(encrypt_text):
    with open("client-private.pem") as f:
        key = f.read()
        print('key=',key)
        rsakey = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsakey)
        text = cipher.decrypt(base64.b64decode(encrypt_text), random_generator)
        print('text=',text)
        return text.decode('utf-8')


def server_encrypt(message, client_public_key):
    rsakey = RSA.importKey(client_public_key)
    cipher = Cipher_pkcs1_v1_5.new(rsakey)
    cipher_text = base64.b64encode(cipher.encrypt(message.encode('utf-8')))
    return cipher_text.decode('utf-8')


def server_decrypt(encrypt_text):
    with open("server-private.pem") as f:
        key = f.read()
        rsakey = RSA.importKey(key)
        cipher = Cipher_pkcs1_v1_5.new(rsakey)
        text = cipher.decrypt(base64.b64decode(encrypt_text), random_generator)
        return text.decode('utf-8')


def client_sign(message):
    with open("client-private.pem") as f:
        key = f.read()
        rsakey = RSA.importKey(key)
        signer = Signature_pkcs1_v1_5.new(rsakey)
        digest = SHA1.new()
        digest.update(bytes(message.encode('utf-8')))
        sign = signer.sign(digest)
        signature = base64.b64encode(sign)
        return signature.decode('utf-8')

def server_verify(message, signature, client_public_key):
    rsakey = RSA.importKey(client_public_key)
    verifier = Signature_pkcs1_v1_5.new(rsakey)
    digest = SHA1.new()
    # Assumes the data is base64 encoded to begin with
    digest.update(bytes(message.encode('utf-8')))
    verify = verifier.verify(digest, base64.b64decode(signature))
    return verify


def server_sign(message):
    with open("server-private.pem") as f:
        key = f.read()
        rsakey = RSA.importKey(key)
        signer = Signature_pkcs1_v1_5.new(rsakey)
        digest = SHA1.new()
        digest.update(bytes(message.encode('utf-8')))
        sign = signer.sign(digest)
        signature = base64.b64encode(sign)
        return signature.decode('utf-8')


def client_verify(message, signature):
    with open("server-public.pem") as f:
        key = f.read()
        rsakey = RSA.importKey(key)
        verifier = Signature_pkcs1_v1_5.new(rsakey)
        digest = SHA1.new()
        # Assumes the data is base64 encoded to begin with
        digest.update(bytes(message.encode('utf-8')))
        verify = verifier.verify(digest, base64.b64decode(signature))
        return verify



def save_pub_key(pub_key):
    with open('server-public.pem', 'w') as f:
        f.write(pub_key)
